import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import numpy as np
from scipy.io import savemat


def load_and_preprocess_data(path):
    # 定义列名
    cols = """duration,
    protocol_type,
    service,
    flag,
    src_bytes,
    dst_bytes,
    land,
    wrong_fragment,
    urgent,
    hot,
    num_failed_logins,
    logged_in,
    num_compromised,
    root_shell,
    su_attempted,
    num_root,
    num_file_creations,
    num_shells,
    num_access_files,
    num_outbound_cmds,
    is_host_login,
    is_guest_login,
    count,
    srv_count,
    serror_rate,
    srv_serror_rate,
    rerror_rate,
    srv_rerror_rate,
    same_srv_rate,
    diff_srv_rate,
    srv_diff_host_rate,
    dst_host_count,
    dst_host_srv_count,
    dst_host_same_srv_rate,
    dst_host_diff_srv_rate,
    dst_host_same_src_port_rate,
    dst_host_srv_diff_host_rate,
    dst_host_serror_rate,
    dst_host_srv_serror_rate,
    dst_host_rerror_rate,
    dst_host_srv_rerror_rate"""

    columns = [c.strip() for c in cols.split(',') if c.strip()]
    columns.append('target')

    # 定义攻击类型映射
    attacks_types = {
        'normal': 'normal',
        'back': 'dos', 'land': 'dos', 'neptune': 'dos', 'pod': 'dos',
        'smurf': 'dos', 'teardrop': 'dos', 'apache2': 'dos', 'udpstorm': 'dos',
        'processtable': 'dos', 'mailbomb': 'dos', 'worm': 'dos',
        'buffer_overflow': 'u2r', 'loadmodule': 'u2r', 'perl': 'u2r',
        'rootkit': 'u2r', 'httptunnel': 'u2r', 'ps': 'u2r', 'sqlattack': 'u2r',
        'xterm': 'u2r',
        'ftp_write': 'r2l', 'guess_passwd': 'r2l', 'imap': 'r2l', 'multihop': 'r2l',
        'phf': 'r2l', 'spy': 'r2l', 'warezclient': 'r2l', 'warezmaster': 'r2l',
        'named': 'r2l', 'sendmail': 'r2l', 'snmpgetattack': 'r2l',
        'snmpguess': 'r2l', 'xlock': 'r2l', 'xsnoop': 'r2l',
        'ipsweep': 'probe', 'nmap': 'probe', 'portsweep': 'probe', 'satan': 'probe',
        'mscan': 'probe', 'saint': 'probe'
    }

    # 读取数据
    df = pd.read_csv(path, names=columns)

    # 添加攻击类型列
    def map_attack(attack):
        attack = attack.strip('.')
        return attacks_types.get(attack, 'unknown')

    df['Attack Type'] = df['target'].apply(map_attack)

    # 删除空值列
    df = df.dropna(axis='columns')

    # # 删除高度相关特征
    # drop_cols = [
    #     'num_root', 'srv_serror_rate', 'srv_rerror_rate',
    #     'dst_host_srv_serror_rate', 'dst_host_serror_rate',
    #     'dst_host_rerror_rate', 'dst_host_srv_rerror_rate',
    #     'dst_host_same_srv_rate'
    # ]
    # df = df.drop([c for c in drop_cols if c in df.columns], axis=1)

    # 特征编码
    pmap = {'icmp': 0, 'tcp': 1, 'udp': 2}
    fmap = {'SF': 0, 'S0': 1, 'REJ': 2, 'RSTR': 3, 'RSTO': 4,
            'SH': 5, 'S1': 6, 'S2': 7, 'RSTOS0': 8, 'S3': 9, 'OTH': 10}

    df['protocol_type'] = df['protocol_type'].map(pmap)
    df['flag'] = df['flag'].map(fmap)

    # 删除不需要的特征
    df = df.drop(['target', 'service'], axis=1, errors='ignore')

    # 分割特征和标签
    y = df[['Attack Type']]
    X = df.drop(['Attack Type'], axis=1)

    # 特征归一化
    scaler = MinMaxScaler()
    X = pd.DataFrame(scaler.fit_transform(X), columns=X.columns)

    return X, y


# 使用示例
if __name__ == "__main__":
    # 原始映射字典
    attack_dict = {
        'normal': 0,
        'dos': 1,
        'u2r': 2,
        'r2l': 3,
        'probe': 4
    }

    path = "./data/corrected.gz"
    X, y = load_and_preprocess_data(path)  # 假设该函数返回两个DataFrame

    # 标签映射
    y_mapped = np.vectorize(attack_dict.get)(y)

    # 将y_mapped转换为DataFrame并重命名列
    y_df = pd.DataFrame(y_mapped, columns=['label'], index=X.index)

    # 拼接数据 (将标签放在第0列)
    concatenated = pd.concat([y_df, X], axis=1)

    # 转换为数值矩阵 (丢弃索引和列名)
    data_matrix = concatenated.values

    # 保存为.mat文件
    savemat('kd99_test.mat',
            {'data': data_matrix,
             'column_names': concatenated.columns.values.tolist()})
